from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate
from operator import itemgetter
from langchain_core.output_parsers import JsonOutputParser
from langchain.tools.render import render_text_description
from langdev_helper.llm.qwen import llm

@tool
def multiply(first_int: int, second_int: int) -> int:
    """将两个整数相乘。"""
    return first_int * second_int

@tool
def add(first_int: int, second_int: int) -> int:
    "将两个整数相加。"
    return first_int + second_int

@tool
def exponentiate(base: int, exponent: int) -> int:
    "对底数求指数幂。"
    return base**exponent

tools = [add, exponentiate, multiply]

def tool_chain(model_output):
    tool_map = {tool.name: tool for tool in tools}
    chosen_tool = tool_map[model_output["name"]]
    return itemgetter("arguments") | chosen_tool

rendered_tools = render_text_description(tools)

DEFAULT_SYSTEM_TEMPLATE = """You have access to the following tools:

{rendered_tools}

You must always select one of the above tools and respond with only a JSON object matching the following schema:

{{
  "tool": <name of the selected tool>,
  "tool_input": <parameters for the selected tool, matching the tool's JSON schema>
}}
"""  # noqa: E501

DEFAULT_RESPONSE_FUNCTION = {
    "name": "__conversational_response",
    "description": (
        "Respond conversationally if no other tools should be called for a given query."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "response": {
                "type": "string",
                "description": "Conversational response to the user.",
            },
        },
        "required": ["response"],
    },
}


# print(DEFAULT_SYSTEM_TEMPLATE)

system_prompt = DEFAULT_SYSTEM_TEMPLATE

prompt = ChatPromptTemplate.from_messages(
    [("system", system_prompt), ("user", "{input}")]
)


chain = prompt | llm | JsonOutputParser() | tool_chain
chain1= prompt | llm | JsonOutputParser()
chain2= prompt | llm


ret = chain1.invoke({"input": "3加3等于", 'rendered_tools': rendered_tools})

print(ret)

